/*
mpicc mpi_test.c -o mpi_test

# 乌镇裸金属测试
LD_PRELOAD=mpitracer.so:/usr/local/openmpi/lib/libmpi.so.40 hipprof --hip-trace --mpi-trace --output-type 1 mpirun  --allow-run-as-root -np 4 ./mpi_test
mpirun  --allow-run-as-root -np 4 -x LD_PRELOAD=mpitracer.so:/usr/local/openmpi/lib/libmpi.so.40 hipprof --hip-trace --mpi-trace --output-type 1 ./mpi_test
hipprof --db-merge .

# AC单节点测试：
mpirun -np 4 -x LD_PRELOAD=mpitracer.so:/public/software/mpi/intelmpi/2021.3.0/lib/release/libmpi.so.12 hipprof --hip-trace --mpi-trace --output-type 1 ./mpi_test

*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mpi.h>

#define MAX_CMD_LEN 256   // 命令最大长度
#define MAX_STATUS_LEN 128 // 状态信息最大长度

int main(int argc, char**argv) {
    MPI_Init(&argc, &argv);

    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);  // 当前进程编号
    MPI_Comm_size(MPI_COMM_WORLD, &size);  // 总进程数

    // 检查是否启动了4个MPI进程
    if (size != 4) {
        if (rank == 0) {
            printf("请启动4个MPI进程：mpirun -np 4 ./mpi_launcher\n");
        }
        MPI_Finalize();
        return 1;
    }

    // 预设参数列表：每个进程对应不同的设备环境变量
    const char* params[] = {
        "HIP_VISIBLE_DEVICES=0",
        "HIP_VISIBLE_DEVICES=1",
        "HIP_VISIBLE_DEVICES=2",
        "HIP_VISIBLE_DEVICES=3"
    };

    // 构造执行test的命令
    char command[MAX_CMD_LEN];
    snprintf(command, sizeof(command), "%s ./MatrixTranspose", params[rank]);

    // --------------------------
    // 阶段1：进程0发送启动信号
    // --------------------------
    if (rank == 0) {
        // 进程0向所有子进程发送准备就绪信号
        int ready_signal = 1;
        for (int i = 1; i < size; i++) {
            MPI_Send(&ready_signal, 1, MPI_INT, i, 0, MPI_COMM_WORLD);
            printf("进程0向进程%d发送启动准备信号\n", i);
        }
    } else {
        // 子进程接收进程0的信号
        int ready_signal;
        MPI_Recv(&ready_signal, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        printf("进程%d收到启动准备信号\n", rank);
    }

    // --------------------------
    // 阶段2：各进程向进程0发送自己的命令
    // --------------------------
    if (rank == 0) {
        // 进程0接收所有子进程的命令
        char recv_cmd[MAX_CMD_LEN];
        for (int i = 1; i < size; i++) {
            MPI_Recv(recv_cmd, MAX_CMD_LEN, MPI_CHAR, i, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            printf("进程0收到进程%d的命令：%s\n", i, recv_cmd);
        }
        // 打印自己的命令
        printf("进程0自己的命令：%s\n", command);
    } else {
        // 子进程向进程0发送自己的命令
        MPI_Send(command, strlen(command)+1, MPI_CHAR, 0, 1, MPI_COMM_WORLD);
        printf("进程%d已向进程0发送命令\n", rank);
    }

    // --------------------------
    // 阶段3：同步所有进程，统一开始执行
    // --------------------------
    MPI_Barrier(MPI_COMM_WORLD);  // 等待所有进程准备就绪
    if (rank == 0) {
        printf("\n所有进程准备就绪，开始执行test程序...\n\n");
    }

    // 执行test程序
    printf("MPI进程%d启动命令：%s\n", rank, command);
    int exit_status = system(command);

    // --------------------------
    // 阶段4：各进程向进程0发送执行结果
    // --------------------------
    if (rank == 0) {
        // 进程0接收所有子进程的执行状态
        char status_msg[MAX_STATUS_LEN];
        for (int i = 1; i < size; i++) {
            MPI_Recv(status_msg, MAX_STATUS_LEN, MPI_CHAR, i, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            printf("进程0收到进程%d的执行状态：%s\n", i, status_msg);
        }
        // 处理自己的执行状态
        if (exit_status == 0) {
            printf("进程0执行成功\n");
        } else {
            printf("进程0执行失败，退出状态：%d\n", exit_status);
        }
    } else {
        // 子进程向进程0发送执行状态
        char status_msg[MAX_STATUS_LEN];
        if (exit_status == 0) {
            snprintf(status_msg, sizeof(status_msg), "执行成功（状态码：%d）", exit_status);
        } else {
            snprintf(status_msg, sizeof(status_msg), "执行失败（状态码：%d）", exit_status);
        }
        MPI_Send(status_msg, strlen(status_msg)+1, MPI_CHAR, 0, 2, MPI_COMM_WORLD);
    }

    MPI_Finalize();
    return 0;
}
